package main

import (
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/rlp"
)

type BranchNode struct {
	Branches [16]Node
	Value    []byte
}

func NewBranchNode() *BranchNode {
	return &BranchNode{
		Branches: [16]Node{},
	}
}

func (b BranchNode) Hash() []byte {
	return crypto.Keccak256(b.Serialize())
}

func (b *BranchNode) SetBranch(nibble Nibble, node Node) {
	b.Branches[int(nibble)] = node
}

func (b *BranchNode) RemoveBranch(nibble Nibble) {
	b.Branches[int(nibble)] = nil
}

func (b *BranchNode) SetValue(value []byte) {
	b.Value = value
}

func (b *BranchNode) RemoveValue() {
	b.Value = nil
}

func (b BranchNode) Raw() []interface{} {
	return b.RawWithDepth(0)
}

func (b BranchNode) RawWithDepth(depth int) []interface{} {
	hashes := make([]interface{}, 17)
	for i := 0; i < 16; i++ {
		if b.Branches[i] == nil {
			hashes[i] = EmptyNodeRaw
		} else {
			node := b.Branches[i]
			
			// 防止过深的递归调用
			if depth > 100 {
				// 如果递归深度过深，直接使用哈希值
				hashes[i] = node.Hash()
				continue
			}
			
			// 根据节点类型决定如何处理
			switch n := node.(type) {
			case *ExtensionNode:
				nodeRaw := n.RawWithDepth(depth + 1)
				nodeSerialized, err := rlp.EncodeToBytes(nodeRaw)
				if err != nil || len(nodeSerialized) >= 32 {
					hashes[i] = node.Hash()
				} else {
					hashes[i] = nodeRaw
				}
			case *BranchNode:
				nodeRaw := n.RawWithDepth(depth + 1)
				nodeSerialized, err := rlp.EncodeToBytes(nodeRaw)
				if err != nil || len(nodeSerialized) >= 32 {
					hashes[i] = node.Hash()
				} else {
					hashes[i] = nodeRaw
				}
			default:
				// 对于 LeafNode 或其他节点类型，直接调用 Raw()
				nodeRaw := node.Raw()
				nodeSerialized, err := rlp.EncodeToBytes(nodeRaw)
				if err != nil || len(nodeSerialized) >= 32 {
					hashes[i] = node.Hash()
				} else {
					// if node can be serialized to less than 32 bits, then
					// use Serialized directly.
					// it has to be ">=", rather than ">",
					// so that when deserialized, the content can be distinguished
					// by length
					hashes[i] = nodeRaw
				}
			}
		}
	}

	hashes[16] = b.Value
	return hashes
}

func (b BranchNode) Serialize() []byte {
	return Serialize(b)
}

func (b BranchNode) HasValue() bool {
	return b.Value != nil
}
